#include<vector>
#include<iostream>
//#include<cstdio>

#include <src/blas_lapack.h>
#include <src/hamiltonian_actor.h>
#include <src/exception.h>
#include <src/numbers.h>
#include <src/indexing.h>
#include <src/lanczos.h>

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::HamiltonianActor -- constructor
//-----------------------------------------------------------------------------------
template <class SCALAR>
eds::HamiltonianActor<SCALAR>::HamiltonianActor(const int ns, const int to, const int cps,
                                                const std::vector<int> & bonds,
                                                const std::vector<double> & bond_strengths,
                                                const std::vector<int> & triangles,
                                                const int nthread)
  : _nthread(nthread),
    _ns(ns),
    _to(to),
    _cps(cps),
    _thresh(1.0e-9),
    _z_array ( eds::z_array_size(ns)    , 0 ),
    _z_lookup( (ns+1)*(ns+1)            , 0 ),
    _offsets ( (to+1)*(ns+1)*(ns+1)*cps , 0 ),
    _spacing ( (to+1)*(ns+1)*cps        , 0 ),
    _work    ( 5*cps + 2                , 0 ),
    _plus_factor(cps, 0.0),
    _minus_factor(cps, 0.0),
    _bonds(bonds),
    _bond_strengths(bond_strengths),
    _triangles(triangles),
    _start_configs(nthread, eds::Configuration(ns, to, cps)),
    _end_cmpds(nthread, -1),
    _thread_work_ptrs(nthread),
    _thread_workspace()
{

  eds::prepare_z_array(ns, &_z_array.at(0), &_z_lookup.at(0));
  eds::prepare_spacing_and_offsets(ns, cps, to, &_offsets.at(0), &_spacing.at(0), &_work.at(0));

  for (int i = 0; i < cps; i++) {
    const double s = 0.5 * (cps-1);
    const double m = i - 0.5 * (cps-1);
    _plus_factor[i]  = std::sqrt( s*(s+1) - m*(m+1) );
    _minus_factor[i] = std::sqrt( s*(s+1) - m*(m-1) );
  }

  // get the number of configurations
  const int nconfig = eds::n_integer_solutions(ns, to, cps-1, &_work.at(0));

  // get the number of configurations per thread
  const int ncpt = nconfig / nthread;

  // get space for threads to store their vectors
  const size_t buffer_size = 1024;
  _thread_workspace.assign( size_t(nthread) * size_t(nconfig) + size_t(nthread+1) * buffer_size, pcps::zero<SCALAR>() );
  for (int i = 0; i < nthread; i++)
    _thread_work_ptrs.at(i) = &_thread_workspace.at( buffer_size * size_t(i+1) + size_t(nconfig) * size_t(i) );

  // get the starting configurations and ending compound indices for each thread
  {
    int count = 0;
    int tnum = 0;
    eds::Configuration config(ns, to, cps);
    _start_configs.at(0) = config;
    for ( ; !config.finished(); config++, count++) {
      if (count == ncpt && tnum < nthread-1) {
        _end_cmpds.at(tnum++) = count;
        _start_configs.at(tnum) = config;
        count = 0;
      }
      //if ( count == tnum*ncpt + ncpt - 1 ) {
      //  if (tnum < 0 || tnum >= nthread)
      //    throw pcps::Exception("bad tnum in ( count == tnum*ncpt + ncpt - 1 ), tunum = %i") % tnum;
      //  _end_cmpds[tnum++] = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
      //                                               &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);
      //}
      //if ( count == nconfig - 1 ) {
      //  _end_cmpds[nthread-1] = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
      //                                                  &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);
      //}
    }
    if (tnum != nthread-1)
      throw pcps::Exception("bad tnum, tunum = %i") % tnum;
    _end_cmpds.at(tnum) = count;
  }

//  _bonds.push_back( 0); _bonds.push_back( 4); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 0); _bonds.push_back( 1); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 1); _bonds.push_back( 4); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 0); _bonds.push_back( 2); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 0); _bonds.push_back( 7); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 2); _bonds.push_back( 7); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 1); _bonds.push_back( 3); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 1); _bonds.push_back( 5); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 3); _bonds.push_back( 5); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 2); _bonds.push_back( 3); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 2); _bonds.push_back( 6); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 3); _bonds.push_back( 6); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 5); _bonds.push_back( 8); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 5); _bonds.push_back(10); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 8); _bonds.push_back(10); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 6); _bonds.push_back( 9); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 6); _bonds.push_back(10); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 9); _bonds.push_back(10); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 7); _bonds.push_back( 9); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 7); _bonds.push_back(11); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 9); _bonds.push_back(11); _bond_strengths.push_back(1.0);
//
//  _bonds.push_back( 4); _bonds.push_back( 8); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 4); _bonds.push_back(11); _bond_strengths.push_back(1.0);
//  _bonds.push_back( 8); _bonds.push_back(11); _bond_strengths.push_back(1.0);

}

template eds::HamiltonianActor<double>::HamiltonianActor(const int, const int, const int,
                                                         const std::vector<int> &,
                                                         const std::vector<double> &,
                                                         const std::vector<int> &,
                                                         const int);
template eds::HamiltonianActor<std::complex<double> >::HamiltonianActor(const int, const int, const int,
                                                                        const std::vector<int> &,
                                                                        const std::vector<double> &,
                                                                        const std::vector<int> &,
                                                                        const int);

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::apply_bond -- operators by a bond
//-----------------------------------------------------------------------------------

template <class SCALAR>
void eds::HamiltonianActor<SCALAR>::apply_bond(eds::Configuration & config,
                                               const int p, const int q, const int curr_cmpd, const SCALAR c, SCALAR * const y,
                                               int * const work) {

  const double mp = config[p] - 0.5 * (_cps-1);
  const double mq = config[q] - 0.5 * (_cps-1);

  // Sz_p Sz_q
  y[curr_cmpd] += mp * mq * c;

  // S+_p S-_q
  config[p]++;
  config[q]--;
  if ( config[p] < _cps && config[q] >= 0 ) {
    const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
                                             &_z_lookup[0], &_z_array[0], config.c_ptr(), work);
    y[cmpd] += 0.5 * _plus_factor[config[p]-1] * _minus_factor[config[q]+1] * c;
  }
  config[p]--;
  config[q]++;

  // S+_q S-_p
  config[q]++;
  config[p]--;
  if ( config[q] < _cps && config[p] >= 0 ) {
    const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
                                             &_z_lookup[0], &_z_array[0], config.c_ptr(), work);
    y[cmpd] += 0.5 * _plus_factor[config[q]-1] * _minus_factor[config[p]+1] * c;
  }
  config[q]--;
  config[p]++;

}

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::operate_by_A -- function to perform y = A d
//-----------------------------------------------------------------------------------

template <class SCALAR>
void eds::HamiltonianActor<SCALAR>::operate_by_A(const int n, SCALAR * const d, SCALAR * const y, const int myrank, const int tid) {

  pcps::xscal(n, pcps::zero<SCALAR>(), y, 1);

  // create thread attribute object to ensure threads are joinable
  pthread_attr_t attr;
  pthread_attr_init(&attr);
  pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);

  // create mutexes
  pthread_mutex_t mtx;
  pthread_mutex_init(&mtx, NULL);

  // start threads
  std::vector<pthread_t> threads(_nthread);
  std::vector<typename eds::HamiltonianActor<SCALAR>::ThreadArgs> thread_args(_nthread);
  for (int i = 0; i < _nthread; i++) {
    thread_args[i].tid = i;
    thread_args[i].obj_ptr = this;
    thread_args[i].mtx_ptr = &mtx;
    thread_args[i].n = n;
    thread_args[i].d = d;
    thread_args[i].y = y;
    thread_args[i].myrank = myrank;
    const int rc = pthread_create(&threads.at(i), &attr, eds::thread_operate_by_A<SCALAR>, (void *)&thread_args[i]);
    if (rc != 0)
      throw pcps::Exception("error in eds::HamiltonianActor::operate_by_A.  pthread_create failed for thread %i with return code %i") % i % rc;
  }

  // wait for threads to finish
  for (int i = 0; i < _nthread; i++) {
    const int rc = pthread_join(threads.at(i), NULL);
    if (rc != 0)
      throw pcps::Exception("error in eds::HamiltonianActor::operate_by_A.  pthread_join failed for thread %i with return code %i") % i % rc;
  }

  // destroy mutexes
  pthread_mutex_destroy(&mtx);

  // destroy thread attribute object
  pthread_attr_destroy(&attr);

}

template void eds::HamiltonianActor<double>::operate_by_A(const int n, double * const d, double * const y, const int myrank, const int tid);
template void eds::HamiltonianActor<std::complex<double> >::operate_by_A(const int n,
                                                                         std::complex<double> * const d,
                                                                         std::complex<double> * const y,
                                                                         const int myrank,
                                                                         const int tid);

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::operate_by_M -- function to perform y = M d
//-----------------------------------------------------------------------------------

template <class SCALAR>
void eds::HamiltonianActor<SCALAR>::operate_by_M(const int n, SCALAR * const d, SCALAR * const y, const int myrank, const int tid) {
  throw pcps::Exception( ( boost::format("function operate_by_M is not implemented for the eds::HamiltonianActor class") ).str() );
}

template void eds::HamiltonianActor<double>::operate_by_M(const int n, double * const d, double * const y, const int myrank, const int tid);
template void eds::HamiltonianActor<std::complex<double> >::operate_by_M(const int n,
                                                                         std::complex<double> * const d,
                                                                         std::complex<double> * const y,
                                                                         const int myrank,
                                                                         const int tid);

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::operate_by_M_inv -- function to perform y = M^(-1) d
//-----------------------------------------------------------------------------------

template <class SCALAR>
void eds::HamiltonianActor<SCALAR>::operate_by_M_inv(const int n, SCALAR * const d, SCALAR * const y, const int myrank, const int tid) {
  throw pcps::Exception( ( boost::format("function operate_by_M_inv is not implemented for the eds::HamiltonianActor class") ).str() );
}

template void eds::HamiltonianActor<double>::operate_by_M_inv(const int n, double * const d, double * const y, const int myrank, const int tid);
template void eds::HamiltonianActor<std::complex<double> >::operate_by_M_inv(const int n,
                                                                             std::complex<double> * const d,
                                                                             std::complex<double> * const y,
                                                                             const int myrank,
                                                                             const int tid);

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::converged -- function to check if the iterations
//                                           have converged
//-----------------------------------------------------------------------------------

template <class SCALAR>
bool eds::HamiltonianActor<SCALAR>::converged(const double residual, const int myrank, const int tid) {

  // ensure we are not threading
  assert( tid == 0 );

  // compute return value on root process
  bool retval;
  if (myrank == 0 && tid == 0) {
    assert( residual >= 0.0 );
    retval = (residual < this->_thresh);
  }

  //// broadcast return value to all processes
  //if (tid == 0) MPI::COMM_WORLD.Bcast(&retval, 1, MPI::BOOL, 0);

  // return whether or not we have converged
  return retval;

}

template bool eds::HamiltonianActor<double>::converged(const double residual, const int myrank, const int tid);
template bool eds::HamiltonianActor<std::complex<double> >::converged(const double residual, const int myrank, const int tid);

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::tp_term -- computes a term used in triple products
//-----------------------------------------------------------------------------------

template <class SCALAR>
SCALAR eds::HamiltonianActor<SCALAR>::tp_term(const int i, const int j, const int k, const SCALAR * const x,
                                              eds::Configuration & config) {

  SCALAR retval = pcps::zero<SCALAR>();

  const double mi = config[i] - 0.5 * (_cps-1);

  // Sz_i S+_j S-_k
  config[j]++;
  config[k]--;
  if ( config[j] < _cps && config[k] >= 0 ) {
    const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
                                             &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);
    retval += mi * _plus_factor[config[j]-1] * _minus_factor[config[k]+1] * pcps::conj(x[cmpd]);
  }
  config[j]--;
  config[k]++;

  // Sz_i S+_k S-_j
  config[k]++;
  config[j]--;
  if ( config[k] < _cps && config[j] >= 0 ) {
    const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
                                             &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);
    retval -= mi * _plus_factor[config[k]-1] * _minus_factor[config[j]+1] * pcps::conj(x[cmpd]);
  }
  config[k]--;
  config[j]++;

  retval *= 0.5 * pcps::imaginary_unity<SCALAR>();

  return retval;

}

//-----------------------------------------------------------------------------------
// eds::HamiltonianActor::triple_products -- computes triangles' triple products
//-----------------------------------------------------------------------------------

template <class SCALAR>
void eds::HamiltonianActor<SCALAR>::triple_products(const int n, const int ntpv, const int tri_num) {

  if ( _triangles.size() % 3 != 0)
    throw pcps::Exception("triangle vector length must be a multiple of three");

  // load solution vectors
  std::vector<std::vector<SCALAR> > sol_vecs(ntpv);
  for (int i = 0; i < ntpv; i++) {
    const std::string dir = "./";
    const std::string solution_prefix = dir + "lanczos_solution_";
    sol_vecs.at(i).assign(n, pcps::zero<SCALAR>());
    pcps::lanczos_load_vec(n, i, solution_prefix, &sol_vecs.at(i).at(0));
  }

  // diagonalize space of solution vectors w.r.t. to triple product operator on the first triangle
  {

    // get triple product and overlap matrices
    std::vector<SCALAR> tp_mat(ntpv*ntpv, pcps::zero<SCALAR>());
    std::vector<SCALAR> s_mat(ntpv*ntpv, pcps::zero<SCALAR>());
    for (int i = 0; i < ntpv; i++) {
      for (int j = i; j < ntpv; j++) {

        // get the overlap elements
        s_mat.at(i*ntpv+j) = pcps::xdotc(n, &sol_vecs.at(i).at(0), 1, &sol_vecs.at(j).at(0), 1);
        s_mat.at(j*ntpv+i) = pcps::conj(s_mat.at(i*ntpv+j));

        // get triple product matrix elements
        tp_mat.at(i*ntpv+j) = pcps::zero<SCALAR>();
        for (eds::Configuration config(_ns, _to, _cps); !config.finished(); config++) {

          // get compound index
          const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
                                                   &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);

          // get coefficient
          const SCALAR coeff = sol_vecs.at(j).at(cmpd);

          const int I = _triangles.at(3*tri_num+0);
          const int J = _triangles.at(3*tri_num+1);
          const int K = _triangles.at(3*tri_num+2);

          SCALAR d = pcps::zero<SCALAR>();
          d += this->tp_term(I, J, K, &sol_vecs.at(i).at(0), config);
          d += this->tp_term(J, K, I, &sol_vecs.at(i).at(0), config);
          d += this->tp_term(K, I, J, &sol_vecs.at(i).at(0), config);
          d *= coeff;

          tp_mat.at(i*ntpv+j) += d;

        }
        tp_mat.at(j*ntpv+i) = pcps::conj(tp_mat.at(i*ntpv+j));

      }
    }

    // diagonalize
    std::vector<double> tp_evals(ntpv, 0.0);
    std::vector<double> s_evals(ntpv, 0.0);
    {

      pcps::xconj(tp_mat.size(), &tp_mat.at(0), 1); // fortran will transpose, so we should also conjugate
      pcps::xconj( s_mat.size(), & s_mat.at(0), 1); // fortran will transpose, so we should also conjugate

      int info = 0;
      int tp_lwork = 10*ntpv;
      std::vector<SCALAR> tp_work(tp_lwork, pcps::zero<SCALAR>());
      std::vector<double> tp_rwork(3*ntpv, 0.0);

      // overlap matrix
      std::vector<SCALAR> s_copy(s_mat);
      pcps::xsyev('V', 'U', ntpv, &s_copy.at(0), ntpv, &s_evals.at(0),
                  &tp_work.at(0), tp_lwork, &tp_rwork.at(0), info);
      if (info != 0)
        throw pcps::Exception("pcps::xsyev failed with error code %i in eds::HamiltonianActor::triple_products") % info;

      // triple product matrix
      pcps::xsygv(1, 'V', 'U', ntpv, &tp_mat.at(0), ntpv, &s_mat.at(0), ntpv, &tp_evals.at(0),
                  &tp_work.at(0), tp_lwork, &tp_rwork.at(0), info);
      if (info != 0)
        throw pcps::Exception("pcps::xsygv failed with error code %i in eds::HamiltonianActor::triple_products") % info;

    }

    // print overlap eigenvalues
    for (int i = 0; i < ntpv; i++)
      std::cout << boost::format(" s eigenvalue %2i = %20.12f") % i % s_evals.at(i) << std::endl;
    std::cout << std::endl;

    // print triple product eigenvalues
    for (int i = 0; i < ntpv; i++)
      std::cout << boost::format("tp eigenvalue %2i = %20.12f") % i % tp_evals.at(i) << std::endl;
    std::cout << std::endl;

  }

//  std::vector<SCALAR> tp_vec(_triangles.size()/3, pcps::zero<SCALAR>());
//
//  // get the normalization constant
//  const SCALAR norm = pcps::xdotc(n, x, 1, x, 1);
//
//  // loop over configurations
//  for (eds::Configuration config(_ns, _to, _cps); !config.finished(); config++) {
//
//    // get compound index
//    const int cmpd = eds::index_from_occ_vec(_ns, _cps, _to, &_offsets[0], &_spacing[0],
//                                             &_z_lookup[0], &_z_array[0], config.c_ptr(), &_work[0]);
//
//    // get coefficient
//    const SCALAR coeff = x[cmpd];
//
//    // loop over triangles
//    for (int p = 0; p < tp_vec.size(); p++) {
//
//      const int i = _triangles.at(3*p+0);
//      const int j = _triangles.at(3*p+1);
//      const int k = _triangles.at(3*p+2);
//
//      SCALAR d = pcps::zero<SCALAR>();
//      d += this->tp_term(i, j, k, x, config);
//      d += this->tp_term(j, k, i, x, config);
//      d += this->tp_term(k, i, j, x, config);
//      d *= coeff / norm;
//
//      tp_vec.at(p) += d;
//
//    }
//
//  }
//
//  // print results
//  std::cout << std::endl;
//  std::cout << boost::format("triple products on the triangles:") << std::endl;
//  for (int p = 0; p < tp_vec.size(); p++) {
//    const int i = _triangles.at(3*p+0);
//    const int j = _triangles.at(3*p+1);
//    const int k = _triangles.at(3*p+2);
//    std::cout << boost::format("  %4i  %4i  %4i  (  %20.12f , %20.12f  )")
//                 % i % j % k % pcps::real(tp_vec.at(p)) % pcps::imag(tp_vec.at(p))
//              << std::endl;
//  }
//  std::cout << std::endl;

}

template void eds::HamiltonianActor<double>::triple_products(const int, const int, const int);
template void eds::HamiltonianActor<std::complex<double> >::triple_products(const int, const int, const int);

template <class SCALAR> void * eds::thread_operate_by_A(void * arg) {

  typedef typename eds::HamiltonianActor<SCALAR>::ThreadArgs * ArgPtrType;

  ArgPtrType ap = (ArgPtrType)arg;

  const int tid = ap->tid;
  const int myrank = ap->myrank;
  eds::HamiltonianActor<SCALAR> & obj = *(ap->obj_ptr);
  pthread_mutex_t * const mtx_ptr = ap->mtx_ptr;
  const int n = ap->n;;
  const SCALAR * const d = ap->d;
  SCALAR * const y = ap->y;

  // get and set to zero this thread's scalar work space
  SCALAR * const my_y = obj._thread_work_ptrs.at(tid);
  pcps::xscal(n, pcps::zero<SCALAR>(), my_y, 1);

  // create a configuration
  eds::Configuration config(obj._ns, obj._to, obj._cps);

  // get integer workspace
  std::vector<int> work( 5 * obj._cps + 2, 0);

  // loop over this thread's configurations
  int count = 0;
  for (config = obj._start_configs.at(tid); true; config++, count++) {

    if (count == obj._end_cmpds[tid]) break;

    // get compound index
    const int cmpd = eds::index_from_occ_vec(obj._ns, obj._cps, obj._to, &obj._offsets[0], &obj._spacing[0],
                                             &obj._z_lookup[0], &obj._z_array[0], config.c_ptr(), &work[0]);

    // get coefficient
    const SCALAR coeff = d[cmpd];

    // apply the bond operators
    for (int i = 0; i < obj._bond_strengths.size(); i++)
      obj.apply_bond(config, obj._bonds.at(2*i), obj._bonds.at(2*i+1), cmpd, obj._bond_strengths.at(i) * coeff, my_y, &work[0]);

  }

  // add the results to the total
  pthread_mutex_lock(mtx_ptr);
  pcps::xaxpy(n, pcps::unity<SCALAR>(), my_y, 1, y, 1);
  pthread_mutex_unlock(mtx_ptr);

}

template void * eds::thread_operate_by_A<double>(void *);
template void * eds::thread_operate_by_A<std::complex<double> >(void *);
